/**
 * 
 */
package edu.umd.clip.lm.programs;

import edu.berkeley.nlp.util.*;

import java.io.*;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.charset.Charset;
import java.util.*;
import java.text.*;
import java.lang.management.*;

import com.sleepycat.je.DatabaseException;
import com.sleepycat.je.Environment;

import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.*;
import edu.umd.clip.jobs.*;
import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.factors.Dictionary;
import edu.umd.clip.lm.storage.AbstractProbTreeStorage;
import edu.umd.clip.lm.storage.BDBProbTreeStorage;
import edu.umd.clip.lm.storage.CachedProbTreeStorage;
import edu.umd.clip.lm.storage.ClientProbTreeStorage;
import edu.umd.clip.lm.storage.remote.ClientStorageConnection;
import edu.umd.clip.lm.util.*;
import edu.umd.clip.lm.nbest.*;
import edu.umd.clip.lm.nbest.NbestFormat.NbestFormatData;


/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class LMNbestListTagger {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-input", required = false, usage = "List of nbest files (Default: stdin)")
		public String input;
        @Option(name = "-output", required = true, usage = "Output directory")
		public String output;
        @Option(name = "-jobs", usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-lm", required = false, usage = "LM ID to train (default: " + LanguageModel.PRIMARY_LM_ID + ")")
		public String lm = LanguageModel.PRIMARY_LM_ID;        
        @Option(name = "-debug", usage = "decoder debug level (default: 0)")
        public int debug = 0;
        @Option(name = "-viterbi", usage = "use Viterbi tagging (default: variational)")
        public boolean viterbi = false;
        @Option(name = "-host", usage = "remote storage hostname or IP address (default: use local storage)")
        public String host = null;
        @Option(name = "-port", usage = "remove storage port (default: 2332)")
        public int port = 2332;
        @Option(name = "-meminfo", usage = "track memory usage (default: false)")
        public boolean meminfo = false;
        @Option(name = "-coarse-threshold", usage = "coarse threshold, requires the new decoder (default: 1e-4)")
        public double coarse = 1e-4;
	}

	/**
	 * @param args
	 * @throws DatabaseException 
	 * @throws IOException 
	 */
	public static void main(String[] args) throws DatabaseException, IOException {
		final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
		
		MemoryUsage memuse;
		
        OptionParser optParser = new OptionParser(Options.class);
        final Options opts = (Options) optParser.parse(args, true);

		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: initial: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

		Experiment.initialize(opts.config);
		Experiment experiment = Experiment.getInstance();

		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after Experiment.initialize(): %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}
		experiment.setDebugDecoder(opts.debug);
		
		JobManager.initialize(opts.jobs);
		Thread thread = new Thread(JobManager.getInstance(), "Job Manager");
		thread.setDaemon(true);
		thread.start();

		experiment.buildPrefixes();
		experiment.buildWordPrefixes();
		
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after building prefixes: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}
		
		final InputParser parser = new PlainInputParser(experiment.getTupleDescription());
		{
			long[] endTuples = new long[1];
			endTuples[0] = experiment.getTupleDescription().createEndTuple(); 
			parser.setEndTuples(endTuples);
		}
		final LanguageModel lm = experiment.getLM(opts.lm);
		
		Environment env = null;
		ClientStorageConnection connection = null;
		Experiment.Files files = experiment.getFiles();
		
		if (opts.host == null) {
			env = BDBProbTreeStorage.createEnvironment(files.getDb(), false);
		} else {
			InetSocketAddress address = new InetSocketAddress(InetAddress.getByName(opts.host), opts.port);
			connection = new ClientStorageConnection(address);
			Timer timer = new Timer(true);
			long interval = ClientStorageConnection.ClientMonitor.INTERVAL;
			timer.scheduleAtFixedRate(new ClientStorageConnection.ClientMonitor(connection), interval, interval);
		}

		{
			BDBProbTreeStorage dbStorage = null;
			InterpolatingProbTreeStorage interpolatingStorage = null;
			ClientProbTreeStorage remoteStorage = null;
			if (opts.host == null) {
				dbStorage = new BDBProbTreeStorage(env);
				interpolatingStorage = new InterpolatingProbTreeStorage(dbStorage);
				interpolatingStorage.setCacheSize(100000);
			} else {
				remoteStorage = new ClientProbTreeStorage(connection);
			}
			// load the chain of LMs
			LanguageModel boLM = lm;
			while(true) {
				boLM.getHistoryTree();
				int hardCacheSize = (int) (3000 * Math.pow(4.0, boLM.getOrder()));
				
				if (opts.host == null) {
					dbStorage.open(boLM.getId(), boLM.getIdNum(), false);
					interpolatingStorage.addLm(boLM);
					boLM.getDecoder().setStorage(interpolatingStorage);
				} else {
					hardCacheSize /= 2;
					boLM.getDecoder().setStorage(remoteStorage);
				}
				if (boLM.getBackoffLM() == null) break;
				boLM = experiment.getLM(boLM.getBackoffLM());
			}
		}
		
		if (opts.coarse > 0) {
			Decoder.USE_COARSE_AND_FINE = true;
			Decoder.COARSE_THRESHOLD = opts.coarse;
		}

		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after loading LMs: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

		final MutableInteger extraWords = new MutableInteger(0);
		if (parser.getStartTuples() != null) {
			extraWords.set(extraWords.intValue() + parser.getStartTuples().length);
		}
		if (parser.getEndTuples() != null) {
			extraWords.set(extraWords.intValue() + parser.getEndTuples().length);
		}

		final DataSequencer<NbestFormatData> sequencer = new DataSequencer<NbestFormatData>(1);
		final Dictionary tagDictionary = experiment.getTupleDescription().getDictionary("T");
		final Tagger tagger = opts.viterbi 
			? new ViterbiTagger(lm.getDecoder(), 10000, "T") 
			: new VariationalTagger(lm.getDecoder(), 1000, "T");

		class EvalRunnable implements Runnable {
			Pair<Integer, NbestFormatData> pair;
			public EvalRunnable(int sentNo, NbestFormatData line) {
				pair = new Pair<Integer, NbestFormatData>(sentNo, line);
			}
			
			public void run() {
		    	long[] sentence = parser.parseSentence(pair.getSecond().getSentence());
		    	ArrayList<String> words = parser.parseSentenceToWords(pair.getSecond().getSentence());
		    	Pair<SortedSet<? extends TaggedPath>,Double> result = tagger.evaluateSentence(sentence);
		    	SortedSet<? extends TaggedPath> nbestTags = result.getFirst();
		    	
		    	//ArrayList<ViterbiPath> bestTags = new ArrayList<ViterbiPath>(nbestTags);
		    	int tags[] = nbestTags.last().getOrderedTags();
				StringBuilder sb = new StringBuilder();
				int firstWord = lm.getOrder();
				for(int i=firstWord-1; i<tags.length-1; ++i) {
					sb.append(words.get(i-firstWord+1));
					sb.append('/');
					sb.append(tagDictionary.getWord(tags[i]));
					sb.append(' ');
				}
		    	pair.getSecond().setSentence(sb.toString());
		    	sequencer.putItem(pair);
				if (opts.meminfo) {
					MemoryUsage memuse = memoryBean.getHeapMemoryUsage();
					System.out.printf("MEMUSE: after evaluating sentence %s: %dM/%dM\n", pair.getFirst().toString(), memuse.getUsed()/1048576, memuse.getMax()/1048576);
					//System.out.println("STATS: " + eval.getStatistics());
				}

			}
		}
		DecoderCompactProbTree emptyTree = DecoderCompactProbTree.constructEmptyTree(Experiment.getInstance().getHFT().getTree());
		DecoderCompactProbTree.setEmptyTree(emptyTree);
		
		experiment.closeXML();
		
		Timer timer = null;
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after experiment.closeXML(): %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
			timer = new Timer(true);
			TimerTask task = new TimerTask() {
				public void run() {
					MemoryUsage memuse = memoryBean.getHeapMemoryUsage();
					DateFormat df = DateFormat.getDateTimeInstance(DateFormat.SHORT, DateFormat.LONG);
					System.out.printf("MEMUSE: %s: %dM/%dM\n", df.format(new Date()), memuse.getUsed()/1048576, memuse.getMax()/1048576);					
				}
			};
			timer.scheduleAtFixedRate(task, 0, 10000);
		}

		try {
			BufferedReader reader = new BufferedReader(new InputStreamReader(opts.input == null ? System.in : new FileInputStream(opts.input), 
					Charset.forName("UTF-8")));

			{
				File outDir = new File(opts.output);
				if (!outDir.exists()) {
					if (!outDir.mkdir()) {
						System.err.printf("failed to create output directory: %s\n", outDir.getPath());
						System.exit(1);
					}
				}
			}
		    JobManager manager = JobManager.getInstance();
		    JobGroup group = manager.createJobGroup("decoding");
		    
		    for(String fname = reader.readLine(); fname != null; fname = reader.readLine()) {
		    	try {
		    		File file = new File(fname);
		    		InputStream in = IO.getInputStream(file);
		    		NbestFile nbest = NbestFile.readFile(in);
		    		int lineNo = 0;
		    		for(NbestFormatData data : nbest.getItems()) {
		    			Job job = new Job(new EvalRunnable(++lineNo, data), "decoding sentence #"+Integer.toString(lineNo));
		    			manager.addJob(group, job);
		    		}
		    		in.close();
		    		group.join();
		    		File outFile = new File(opts.output, file.getName());
		    		OutputStream out = IO.getOutputStream(outFile);
		    		nbest.writeFile(out);
		    		out.close();
		    		System.out.printf("done rescoring %s\n", file.getName());
		    	} catch(Exception e) {
		    		System.out.printf("Error processing file %s\n", fname);
		    		e.printStackTrace(System.out);
		    	}
		    }
		} catch(IOException e) {
			e.printStackTrace();
		}
	}

}
